/*
 * EquilibriumSolver.h
 *
 *  Created on: 13 Jun 2011
 *      Author: allan
 */

#ifndef EQUILIBRIUMSOLVER_H_
#define EQUILIBRIUMSOLVER_H_

// C++ includes
#include <vector>
using namespace std;

// Eigen includes
#include <Eigen/Core>
using namespace Eigen;

// GeoReact includes
#include "Functions.h"
#include "MultirootSolver.h"
#include "Reaction.h"
#include "System.h"

class EquilibriumProblem;

class EquilibriumSolver
{
public:
	EquilibriumSolver();
	
	EquilibriumSolver(const System& system, const StringVector& equilibriumSpecies);
	
	const bool
	SpecifySpeciesTotalMoles(const String& speciesName, double val);
	
	const bool
	SpecifySpeciesMoles(const String& speciesName, double val);
	
	const bool
	SpecifySpeciesActivity(const String& speciesName, double val);
	
	const bool
	SpecifyAcidity(double pH);
	
	const bool
	SpecifyGasPartialPressure(const String& gasName, double Pg);
	
	const bool
	SpecifyChargeBalance();
	
	void
	ResetConditions();
	
	void
	SetIndependentVector(const DoubleVector& h);
	
	void
	SetSolverOptions(double tolerance, uint maxiterations);
	
	const DoubleVector
	Solve(double T, double P, const DoubleVector& n0);
	
	friend class EquilibriumProblem;
	
private:
	/// The system object which is used to compute activities of the species for every phase 
	System system;
	
	/// The equilibrium-controlled reactions that defines the set of mass-action law equations
	vector<ReactionPtr> reactions;
	
	/// The number of primary (Nj) and equilibrium (Ne) species
	uint Nj, Ne;
	
	/// The indexes of the primary (jSpecies) and equilibrium (eSpecies) species
	IndexVector jSpecies, eSpecies;
	
	/// The stoichiometric matrix of the provided equilibrium-controlled reactions
	MatrixXd vej;
	
	/// The matrices that define the set of equilibrium conditions for the primary and/or equilibrium species
	MatrixXd alpha, beta, theta, mu;
	
	/// The functions that evaluate the right-hand side vector h of the equations for the equilibrium conditions
	vector<SysStateFunction> hfunctions;
	
	/// The tolerance for the iterative procedure
	double tolerance;
	
	/// The maximum number of iterations for the iterative procedure
	uint maxiterations;
	
	/// The current number of specified equilibrium constraints of the solver (this value must equal Nj before the solver is executed)
	uint numConditions;
};

class EquilibriumProblem : public MultirootProblem
{
public:
	EquilibriumProblem(const EquilibriumSolver& equilibriumSolver, double T, double P, const DoubleVector& n0);
	
	virtual void Function(const VectorXd& njne, VectorXd& F);
	
	virtual void Jacobian(const VectorXd& njne, MatrixXd& J);
	
	virtual void PostProcess(VectorXd& njne);
	
	void UpdateAuxiliarData(const VectorXd& njne);
	
	void UnderRelax(VectorXd& njne, const VectorXd& njneOld);
	
private:
	/// A reference to the equilibrium solver object to use the data already created there (note that the class EquilibriumSolver is a friend of EquilibriumProblem)
	const EquilibriumSolver& es;
	
	/// The temperature (T), pressure (P), and initial-guess mole vector (n0)
	const double T, P;
	const DoubleVector& n0;
	
	/// The number of primary (Nj) and equilibrium (Ne) species
	uint Nj, Ne;
	
	/// The mole vector of all species (n), of the primary species (nj), and equilibrium species (ne)  
	DoubleVector n;	
	VectorXd nj, ne;
	
	/// The activity vector of all species (a), activity vector of the primary species (aj), and activity vector of the equilibrium species (ae)
	DoubleVector a;
	VectorXd aj, ae;
	
	/// The equilibrium constants (Ke) and the reaction quotients (Qe) of the equilibrium-controlled reactions
	VectorXd Ke, Qe;
	
	/// The right-hand side vector h of the equations for the equilibrium conditions
	VectorXd h;
	
	/// The primary (Fj) and equilibrium (Fe) part of the residual vector F
	VectorXd Fj, Fe;
	
	/// The block matrices that define the jacobian of the linearised system of non-linear equations 
	MatrixXd Jjj, Jje, Jej, Jee;
	
	/// The solution vector (njne) at the previous iteration (used for application of under relaxation in the solution)
	VectorXd njneOld;
};


//class EquilibriumConstraint
//{
//public:
//	EquilibriumConstraint();
//	
//	EquilibriumConstraint(const EquilibriumSolver& equilibriumSolver);
//	
//	const bool
//	SpecifySpeciesTotalMoles(const String& speciesName, double val);
//	
//	const bool
//	SpecifySpeciesMoles(const String& speciesName, double val);
//	
//	const bool
//	SpecifySpeciesActivity(const String& speciesName, double val);
//	
//	const bool
//	SpecifyAcidity(double pH);
//	
//	const bool
//	SpecifyGasPartialPressure(const String& gasName, double Pg);
//	
//	const bool
//	SpecifyChargeBalance();
//	
//private:
//	/// The system which this equilibrium constraint object is based on
//	const System* system;
//	
//	/// The indexes of the primary species (jIndexes) and equilibrium species (eIndexes) in the system
//	IndexVector jIndexes, eIndexes;
//	
//	/// The matrices that define the set of equilibrium constraints in the system 
//	MatrixXd alpha, beta, theta, mu;
//	
//	/// The functions that evaluate the right-hand side vector h of the system of equilibrium-constraint equations
//	vector<SysStateFunction> hfunctions;
//};

#endif /* EQUILIBRIUMSOLVER_H_ */
